open Ext
open Printf
open BnType

(*
 * Read a file in CN format
 *)

let of_parse (schema, (acyclic, varfactors)) = 

  let numvars = Array.length schema in
  let bn = create_empty_network schema in

  (* Create arrays for distributions and parent lists.
   * We'll fill them in later.
   *)
  let dists = Array.create numvars (Table [||]) in
  let parents = Array.create numvars [] in

  (* Construct CPTs *)
  let build_dist (var, pfactors) =
    let factors = List.map (Mn.Factor.pfactor_to_factor schema) pfactors in
    dists.(var) <- FactorSet factors;
    parents.(var) <- factorset_parents numvars var factors in

  (* This builds all of the CPTs, places them in the dists array,
   * and fills in the parents array. *)
  List.iter build_dist varfactors;
  let children = make_children parents in
  let topo_vars = make_topo_vars bn.vars parents children in
  
  (* Put it all together in the final BN *)
  {bn with acyclic=acyclic; dists=dists; topo_vars=topo_vars; 
   parents=parents; children=children}


let parse channel =
  let schema = Data.input_example channel in
  let lexbuf = Lexing.from_channel channel in
  (schema, CnParser.cn CnLexer.lexer lexbuf)


let load channel = of_parse (parse channel)


(*
 * Print a file in CN format
 *)

let output out cn =
  Data.output_schema out (BnType.schema cn);
  if cn.acyclic then
    output_string out "BN {\n"
  else
    output_string out "DN {\n";
  for v = 0 to numvars cn - 1 do
    fprintf out "v%d {\n" v;
    let fl = cpd_to_factors cn v cn.dists.(v) in
    List.iter (Mn.Factor.output_factor out) fl;
    fprintf out "}\n";
  done;
  output_string out "}\n"
